package com.leetcode.tree.bst;

import com.leetcode.basic.TreeNode;

import java.util.PriorityQueue;
import java.util.Queue;

/**
 * @author Dennis Li
 * @date 2020/7/10 22:23
 */
public class KthSmallest_230 {

    Queue<Integer> nodes;
    private int k;

    public int kthSmallest(TreeNode root, int k) {
        this.k = k;
        // 大小为k的优先队列
        nodes = new PriorityQueue<>((o1, o2) -> o2 - o1);
        preOrder(root);
        return nodes.peek();
    }

    public void preOrder(TreeNode root) {
        if (root == null) return;
        enqueue(root);
        preOrder(root.left);
        preOrder(root.right);
    }

    public void enqueue(TreeNode node) {
        if (nodes.size() == k && nodes.peek() > node.val) {
            nodes.poll();
            nodes.offer(node.val);
        } else if (nodes.size() < k) {
            nodes.offer(node.val);
        }
    }

    int ans;
    int count = 0;
    public int kthSmallest3(TreeNode root, int k) {
        inorder(root, k);
        return ans;
    }

    public void inorder(TreeNode root, int k) {
        if (root == null) return;
        inorder(root.left, k);
        count++;
        if(count == k) {
            ans = root.val;
            return;
        }
        inorder(root.right, k);
    }


    public int kthSmallest2(TreeNode root, int k) {
        // 大小为k的优先队列
        // 因为左边的始终是较小值，计算左边还剩下多少个数
        int leftCount = count(root.left);
        // 恰好等于k - 1，代表此时根节点即为第k小的节点
        if (leftCount == k - 1) return root.val;
        // 如果左侧比k-1大，那么证明第k大的点在左子树
        if (leftCount > k - 1) return kthSmallest2(root.left, k);
        // 否则在右子树中，右子树里需要寻找的点就是k-leftCount-1大的点
        return kthSmallest2(root.right, k - leftCount - 1);
    }

    // 记录BST的剩余结点数
    private int count(TreeNode root) {
        if (root == null) return 0;
        return 1 + count(root.left) + count(root.right);
    }

}
